#!/usr/bin/perl -wc

package RFCMail::DNS::AsyncQuery;

use strict;
use base qw(Exporter);
use Carp;
use Errno;
use IO::Socket::INET;
use Event;
use Net::DNS;
use RFCMail::Util::Callback qw(callback);


BEGIN {
    our @EXPORT_OK = qw(async_dns_query
			async_dns_multi_query
			async_dns_query_udp
			async_dns_query_tcp
			async_dns_query_existing_tcp);
}


INIT {
    our $DEFAULT_PORT = 53; # the standard DNS port
    our $DEFAULT_TCP_TIMEOUT = 120;
}


sub async_dns_query {
    my ($callback, @query) = @_;
    async_dns_query_udp([bless([$callback, \@query]), '_udp_query_result'], @query);
}
sub _udp_query_result {
    my ($self, $result) = @_;
    my ($callback, $query) = @$self;
    if ($result and $result->header->tc) {
	# Response was truncated; try TCP on same server
	async_dns_query_tcp($callback, $result->answerfrom, @$query);
    }
    else {
	callback($callback, $result);
    }
}


sub async_dns_multi_query {
    my ($callback, @queries) = @_;
    my $handler = bless {cb => $callback, count => scalar(@queries), results => []};
    for (@queries) { async_dns_query([$handler, '_query_result'], @$_) }
}
sub _query_result {
    my ($self, $result) = @_;
    push @{$self->{results}}, $result
	if $result;
    --$self->{count}
	or callback($self->{cb}, $self->{results});
}


sub async_dns_query_udp {
    my $callback = shift;
    my $query = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    RFCMail::DNS::AsyncQuery::UDP::start($callback, $query);
}


sub async_dns_query_tcp {
    my $callback = shift;
    my $ns = shift;
    our ($DEFAULT_PORT);
    my $sock = IO::Socket::INET->new(PeerAddr => $ns,
				     PeerPort => $DEFAULT_PORT,
				     Proto    => 'tcp',
				     Blocking => 0)
	or return callback($callback, undef);
    my $handler = bless {cb => $callback, sock => $sock};
    async_dns_query_existing_tcp([$handler, '_tcp_query_result'], $sock, @_);
}
sub _tcp_query_result {
    my ($self, $result) = @_;
    $self->{sock}->close;
    # discard result if it is an exception other than "NXDOMAIN"
    if ($result) {
	my $rcode = $result->header->rcode;
	undef $result
	    if $rcode ne 'NOERROR' and $rcode ne 'NXDOMAIN';
    }
    callback($self->{cb}, $result);
}


sub async_dns_query_existing_tcp {
    my $callback = shift;
    my $sock = shift;
    our ($DEFAULT_TCP_TIMEOUT);
    my $packet = Net::DNS::Packet->new(@_)
	or croak("unable to generate query packet");
    my $data = pack("n/a*", $packet->data);
    my $self = bless {cb => $callback, sock => $sock, data => $data};
    Event->io(fd      => $sock,
	      poll    => 'w',
	      timeout => $DEFAULT_TCP_TIMEOUT,
	      cb      => [$self, '_send_tcp_request'],
	     );
}
sub _send_tcp_request {
    my ($self, $event) = @_;
    our ($DEFAULT_TCP_TIMEOUT);
    if ($event->got) {
	# ready to write
	my $wrote = $self->{sock}->syswrite($self->{data});
	if (defined($wrote)) {
	    # cut the bit we sent from $data
	    substr($self->{data}, 0, $wrote, "");
	    # return if more to send
	    return
		if $self->{data};
	}
	else {
	    # some kind of error; just return if it's a blocking issue
	    return
		if $!{EAGAIN} || $!{EWOULDBLOCK};
	}
    }
    # reaching this point means a timeout, a fatal error, or successful transmission
    # no more write-watcher events in any case
    $event->w->cancel;
    # untransmitted data means we failed
    if ($self->{data}) {
	callback($self->{cb}, undef);
	return;
    }
    # success, on the other hand, means we read back the response
    # (response is read into $self->{data}, which we have just proved to be empty)
    Event->io(fd      => $self->{sock},
	      poll    => 'r',
	      timeout => $DEFAULT_TCP_TIMEOUT,
	      cb      => [$self, '_read_tcp_response'],
	     );
}
sub _read_tcp_response {
    my ($self, $event) = @_;
    my $packet;
    if ($event->got) {
	# ready to read; compute the remaining bytes wanted
	# note that we first want two bytes to determine the size, then that many more
	my $have = length($self->{data});
	my $want = ($have < 2 ? 2 : unpack("n", $self->{data}) + 2) - $have;
	my $got = $self->{sock}->sysread($self->{data}, $want, $have);
	if (defined($got)) {
	    # some data has been appended; return if not complete
	    return
		if length($self->{data}) < 2;
	    return
		if length($self->{data}) < unpack("n", $self->{data}) + 2;
	    # we seem to have a winner; extract the packet data
	    my $data = substr($self->{data}, 2);
	    $packet = Net::DNS::Packet->new(\$data);
	}
	else {
	    # some kind of error; just return if it's a blocking issue
	    return
		if $!{EAGAIN} || $!{EWOULDBLOCK};
	}
    }
    # reaching this point means a timeout, a fatal error, or successful reception
    # these are all terminal cases, and $packet is correct for all cases
    $event->w->cancel;
    callback($self->{cb}, $packet);
}



# This sub-package handles transmission of a specific UDP query-in-progress.
# It is meant for internal use only, and should be accessed through wrappers.

package RFCMail::DNS::AsyncQuery::UDP;

use strict;
use IO::Socket::INET;
use Net::DNS;
use Event;
use RFCMail::Util::Callback qw(callback);


INIT {
    our @DEFAULT_NAMESERVERS = Net::DNS::Resolver->new->nameservers;
    our $UDP_BUFFER_SIZE = 1500; # technically, the payload shouldn't exceed 512
    our $DEFAULT_PORT = $RFCMail::DNS::AsyncQuery::DEFAULT_PORT;
    our $DEFAULT_MAXTRANS = 4; # maximum number of transmission attempts
    our $DEFAULT_WAIT = 5; # number of seconds to wait between transmission attemtps
    our (%Blackboard, $Socket, $Watcher); # for queries in progress
}


# initiate a query
sub start {
    my ($callback, $query, %opts) = @_;
    our (@DEFAULT_NAMESERVERS, $DEFAULT_PORT, $DEFAULT_MAXTRANS, $DEFAULT_WAIT);
    our (%Blackboard, $Socket, $Watcher);
    # change query ID if it is already on the blackboard
    my $id = $query->header->id;
    while (exists($Blackboard{$id})) { ++$id; $id %= 0x10000; }
    $query->header->id($id);
    # build query-in-progress object; set default values as needs be
    my $self = bless \%opts;
    $self->{callback}      = $callback;
    $self->{query}         = $query;
    $self->{nameservers} ||= [ @DEFAULT_NAMESERVERS ];
    $self->{port}        ||= $DEFAULT_PORT;
    $self->{maxtrans}    ||= $DEFAULT_MAXTRANS;
    $self->{wait}        ||= $DEFAULT_WAIT;
    if (not scalar(%Blackboard)) {
	# this will be the first item on the blackboard; create the socket and watcher
	$Socket = IO::Socket::INET->new(Proto => 'udp', Blocking => 0)
	    or die "unable to open socket for DNS queries ($!)";
	$Watcher = Event->io(fd => $Socket, cb => \&read_response);
    }
    $Blackboard{$id} = $self;
    $self->{sender} = Event->timer(after    => 0,
				   interval => $self->{wait},
				   cb       => [ $self, 'transmit_query' ]);
}


# handle an incoming response on the socket (Event->io driven)
sub read_response {
    my $event = shift;
    my ($data, $from, $answer, $self, $question);
    our ($UDP_BUFFER_SIZE, %Blackboard);
    # receive incoming packet, aborting if any error is encountered
    $from = $event->w->fd->recv($data, $UDP_BUFFER_SIZE, MSG_DONTWAIT)
	or return;
    # convert to dotted-quad format
    $from = inet_ntoa((unpack_sockaddr_in($from))[1]);
    # abort if data doesn't make sense as a DNS message
    $answer = Net::DNS::Packet->new(\$data)
	or return;
    $answer->answerfrom($from);
    # abort if the ID is not listed on the blackboard
    $self = ($Blackboard{$answer->header->id})
	or return;
    # it's a real, expected response, but is it what we asked? abort if not
    $question = ($self->{query}->question)[0]->string;
    for ($answer->question) {
	return if $_->string ne $question;
    }
    if ($answer->header->rcode eq 'NOERROR' or $answer->header->rcode eq 'NXDOMAIN') {
	# successful query: finish up
	$self->finish($answer);
    }
    else {
	# an error! don't ask this unhelpful nameserver again: remove it from the list
	@{$self->{nameservers}} = grep($_ ne $from, @{$self->{nameservers}});
	# coerce a retransmit timeout if that was in response to our most recent transmission
	$self->{sender}->again
	    if $from eq $self->{last_ns};
    }
}


# method to attempt a transmission of the query (or force failure by timeout)
sub transmit_query {
    my ($self) = @_;
    our ($Socket);
    if ($self->{maxtrans} and @{$self->{nameservers}}) {
	--$self->{maxtrans};
	# cycle through nameservers on successive queries
	my $ns = shift @{$self->{nameservers}};
	push @{$self->{nameservers}}, $ns;
	# note that this was our most recent nameserver for optimisation purposes
	$self->{last_ns} = $ns;
	my $ns_addr = pack_sockaddr_in($self->{port}, inet_aton($ns));
	# attempt a send; if it works, it works
	$Socket->send($self->{query}->data, MSG_DONTWAIT, $ns_addr);
    }
    else {
	# we're out of bullets (or targets), so give up
	$self->finish(undef);
    }
}


# method to terminate a query, successfully or not
sub finish {
    my ($self, $response) = @_;
    our (%Blackboard, $Socket, $Watcher);
    # dismantle the sender event generator
    (delete $self->{sender})->cancel;
    # wipe ourself off the blackboard
    delete $Blackboard{$self->{query}->header->id};
    if (not scalar(%Blackboard)) {
	# the blackboard is clear; cancel the watcher and close the socket
	$Watcher->cancel;
	undef $Watcher;
	$Socket->close;
	undef $Socket;
    }
    callback($self->{callback}, $response);
}

1;
